package com.chukun.multi.pipeline.reusable;

import com.chukun.multi.twophase.TerminationToken;

import java.util.Objects;
import java.util.concurrent.*;

/**
 * 基于线程池的Pipe实现类
 * 提交到该Pipe的任务由指定个数的工作线程共同处理
 * @param <IN>
 * @param <OUT>
 */
public class ThreadPoolPipeDecorator<IN,OUT> implements Pipe<IN,OUT> {

    private final Pipe<IN,OUT> delegate;
    private final ExecutorService executorService;
    //线程停止标志
    private final PipeTerminationToken terminationToken;
    private final CountDownLatch stageProcessDownLatch = new CountDownLatch(1);

    public ThreadPoolPipeDecorator(Pipe<IN,OUT> delegate, ExecutorService executorService) {
        this.delegate = delegate;
        this.executorService = executorService;
        this.terminationToken = PipeTerminationToken.newInstance(executorService);
    }


    @Override
    public void setNextPipe(Pipe<IN, OUT> nextPipe) {
        delegate.setNextPipe(nextPipe);
    }

    @Override
    public void init(PipeContext pipeCtx) {
        delegate.init(pipeCtx);
    }

    @Override
    public void shutdown(long timeout, TimeUnit unit) {
        terminationToken.setIsShutdown();
        if (terminationToken.reservations.get()>0) {
            try {
                if (stageProcessDownLatch.getCount() > 0) {
                    stageProcessDownLatch.await(timeout, unit);
                }
            }catch (InterruptedException e){

            }
        }
        delegate.shutdown(timeout,unit);
    }

    @Override
    public void process(IN input) throws InterruptedException {
        Runnable r = new Runnable() {
            @Override
            public void run() {
                int remainingReservations = -1;
                try{
                    delegate.process(input);
                }catch (InterruptedException e) {

                }finally {
                    remainingReservations = terminationToken.reservations.decrementAndGet();
                }
                if (terminationToken.isToShutdown() && remainingReservations==0) {
                    stageProcessDownLatch.countDown();
                }
            }
        };
        executorService.submit(r);
        terminationToken.reservations.incrementAndGet();
    }

    private static class PipeTerminationToken extends TerminationToken{
        private final static ConcurrentMap<ExecutorService,PipeTerminationToken> INSTANCE_MAP  = new ConcurrentHashMap<>();

        private PipeTerminationToken(){}

        void setIsShutdown() {
            this.toShutdown = true;
        }

        static PipeTerminationToken newInstance(ExecutorService executorService) {
            PipeTerminationToken token = INSTANCE_MAP.get(executorService);
            if (Objects.isNull(token)) {
                token = new PipeTerminationToken();
                PipeTerminationToken existToken = INSTANCE_MAP.putIfAbsent(executorService,token);
                if (!Objects.isNull(existToken)) {
                    token = existToken;
                }
            }
            return token;
        }
    }
}
